#include "layers/rope/rope.h"

#include "operators/rope/operator_description.h"
#include "operators/rope/invoke_params.h"
#include "operators/rope/operator.h"
#include "operators/names.h"
#include "operators/operator_container.h"

namespace txdnn {

txdnnStatus_t RopeForward(Handle& handle,
                          const TensorDescriptor& xDesc,
                          const void* x,
                          const TensorDescriptor& cosDesc,
                          const void* cos,
                          const TensorDescriptor& sinDesc,
                          const void* sin,
                          const TensorDescriptor& yDesc,
                          void *y) {

  const auto op_desc  = operators::rope::OperatorDescriptionFwd(xDesc, cosDesc, sinDesc, yDesc);
  const auto& name    = operators::AlgorithmName{"RoPEForward"};
  const auto params   = [&]() {
    auto rope_params      = operators::rope::FwdInvokeParams{};
    rope_params.type      = InvokeType::Run;
    rope_params.x_desc_   = &xDesc;
    rope_params.cos_desc_ = &cosDesc;
    rope_params.sin_desc_ = &sinDesc;
    rope_params.y_desc_   = &yDesc;
    rope_params.x_        = x;
    rope_params.cos_      = cos;
    rope_params.sin_      = sin;
    rope_params.y_        = y;
    return rope_params;
  }();

  const auto& op      = operators::OperatorContainer<operators::rope::RoPEForward>{};

  op.ExecutePrimitive(handle, op_desc, name, params);
  return txdnnStatus_t::TXDNN_STATUS_SUCCESS;
}

} // namespace txdnn